'==================================================================================================
'	Author		David Segura
'	URL			http://winpeguy.com/
'	Source		http://winpe.codeplex.com/
'	Version		20130910
'==================================================================================================
'	Purpose:
'	To handle customizations for WinPE by replacing the default Unattend.xml
'==================================================================================================
'	File Locations in Windows:
'	%DeployRoot%\Boot\+WinPE\Menu\*
'==================================================================================================
'	File Location in WinPE:
'	X:\Unattend.vbs
'	X:\Unattend.xml
'	X:\Unattend_%TEMPLATE%.xml		%TEMPLATE% = "LiteTouchPE" or "Generic"
'==================================================================================================
'	Configuration:
'	None.  Editing this file is not recommended
'==================================================================================================
'	Change Log:
'	20130910	Codeplex Beta
'	20130909	Modified Marker
'	20130905	Modified location of Unattend files to the Menu
'	20130904	Modified location of +WinPE.
'	20130903	Removed Readme.txt and edited Templates to reflect this modification
'	20130830	Removed Expire for PXE
'	20130829	Initial Codeplex Release
'	20130827	Added LiteTouchPE / Generic support (Menu disabled for LiteTouchPE) through Unattend_%TEMPLATE%.txt
'	20130822	Initial Release
'				Determines method of WinPE Boot (CD-ROM\PXE\USB) and writes to HKEY_LOCAL_MACHINE\SYSTEM\@WinPE\BootMediaType
'				Determines WinPE Boot Drive Letter and writes to HKEY_LOCAL_MACHINE\SYSTEM\@WinPE\BootMediaDrive
'				If BootDrive is a USB and file X:\Windows\System32\dpinst.exe exists, the following command will be executed
'					dpinst /S /PATH BootDrive:\+Drivers\Boot
'				If BootDrive is a USB and file BootDrive:\+Update\WinPE.vbs exists, the following command will be executed
'					BootDrive:\+Update\WinPE.vbs
'				If X:\Windows\System32\WindowsPowershell\v1.0\powershell.exe exists, the following command will be executed
'					powershell Set-ExecutionPolicy Bypass -Force
'				If the following files exist then the X:\Menu\LaunchMenu.vbs will be launched using a Process Monitor to prevent WinPE automatic reboot
'					X:\Menu.txt
'					X:\Menu\LaunchMenu.vbs
'					X:\Menu\Menu.hta
'				If the X:\Menu\LaunchMenu.vbs was not launched, and X:\Windows\System32\BGInfo.exe and BGInfo.bgi exists, then BGInfo will be launched
'				If the X:\Menu\LaunchMenu.vbs was not launched, then X:\Deploy\Scripts\LiteTouch.wsf will be launched using wscript.exe
'==================================================================================================
Option Explicit

	Dim objFSO, WshShell, WshNetwork, objShell, sCmd
	Set objFSO					= CreateObject("Scripting.FileSystemObject")
	Set WshShell				= WScript.CreateObject("WScript.Shell")
	Set WshNetwork				= WScript.CreateObject("WScript.Network")
	Set objShell				= WScript.CreateObject("WScript.Shell")

	Dim strComputer, objWMIService
	strComputer					= "."
	Set objWMIService			= GetObject("winmgmts:\\" & strComputer & "\root\CIMV2")

	Dim objRegistry, strKeyPath, strValueName, strValue
	Set objRegistry				= GetObject("winmgmts:\\" & strComputer & "\root\default:StdRegProv")
	Const HKEY_CLASSES_ROOT		= &H80000000
	Const HKEY_CURRENT_USER		= &H80000001
	Const HKEY_LOCAL_MACHINE	= &H80000002
	Const HKEY_USERS			= &H80000003
	Const HKEY_CURRENT_CONFIG	= &H80000005
'==================================================================================================
'Function to check if a Registry Key or String exists
'For example if I want to check if 'HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Control\PEBootType' exists, then I would do:
'If RegKeyExists("HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Control\PEBootType") Then Wscript.Echo "It Exists!"
	Function RegKeyExists(Key)
		Dim Entry
		On Error Resume Next
		Entry = objShell.RegRead(Key)
		If Err.Number <> 0 Then
			Err.Clear
			RegKeyExists = False
		Else
			Err.Clear
			RegKeyExists = True
		End If
	End Function
'==================================================================================================
'Gather Environment Variables and save them into variables we can use in this script 
	Dim envComputerName, envProcessorArchitecture, envSystemDrive, envUserName, envSystemRoot, envTemp, envSystem32, strMenuDir
	envComputerName				= WshShell.ExpandEnvironmentStrings("%ComputerName%")
	envProcessorArchitecture	= WshShell.ExpandEnvironmentStrings("%PROCESSOR_ARCHITECTURE%")	'Will return x86 or AMD64
	envSystemDrive				= WshShell.ExpandEnvironmentStrings("%SystemDrive%")
	envUserName					= WshShell.ExpandEnvironmentStrings("%UserName%")
	envSystemRoot				= WshShell.ExpandEnvironmentStrings("%Windir%")
	envTemp						= WshShell.ExpandEnvironmentStrings("%Temp%")
	envSystem32					= envSystemRoot & "\System32"
	strMenuDir					= envSystemDrive & "\Menu"
'==================================================================================================
'Set PowerShell Execution Policy
	If objFSO.FileExists(envSystem32 & "\WindowsPowershell\v1.0\powershell.exe") Then
		sCmd = "powershell Set-ExecutionPolicy Bypass -Force"
		objShell.Run sCmd, 2, True
	End If
'==================================================================================================
'Determine how WinPE was booted
	Dim strPEBootType, strPEBootDrive
	strKeyPath = "SYSTEM\CurrentControlSet\Control"
	strValueName = "PEBootType"
	objRegistry.GetStringValue HKEY_LOCAL_MACHINE,strKeyPath,strValueName,strPEBootType
	'This will return one of the following values
	'PEBootType = Error
	'PEBootType = Flat
	'PEBootType = Remote						(PXE)
	'PEBootType = Ramdisk:SourceIdentified		(USB)
	'PEBootType = Ramdisk:SourceUnidentified
	'PEBootType = Ramdisk:OpticalDrive			(CD-ROM)
'==================================================================================================
'Write the BootMediaType to the Registry and determine if there is a Drive Letter installed
	strKeyPath = "SYSTEM\@WinPE"

	If strPEBootType = "Remote" Then
		strPEBootType = "PXE"
		objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaType",strPEBootType
		objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaDrive","NONE"
		
	ElseIf strPEBootType = "Ramdisk:SourceIdentified" Then 
		strPEBootType = "DRIVE"
		objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaType",strPEBootType
		objRegistry.GetStringValue HKEY_LOCAL_MACHINE,"SYSTEM\CurrentControlSet\Control","PEBootRamdiskSourceDrive",strPEBootDrive
		objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaDrive",strPEBootDrive
		
		'Install WinPE Drivers that are used on Boot
		If objFSO.FileExists(envSystemDrive & "\Windows\System32\dpinst.exe") and objFSO.FolderExists(strPEBootDrive & "+Drivers\Boot") Then
			sCmd = "dpinst /S /PATH " & strPEBootDrive & "+Drivers\Boot"
			objShell.Run sCmd, 1, True
		End If
		
		'Check for strPEBootDrive:\+Update\WinPE.vbs and run if present
		If objFSO.FileExists(strPEBootDrive & "+Update\WinPE.vbs") Then
			sCmd = strPEBootDrive & "+Update\WinPE.vbs"
			objShell.Run sCmd, 1, True
		End If
		
	ElseIf strPEBootType = "Ramdisk:OpticalDrive" Then 
		strPEBootType = "CD-ROM"
		objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaType",strPEBootType
		
		'Need to determine the drive used to boot
		Dim colCD, objCD, colItems
		Set colCD = objWMIService.ExecQuery("SELECT * FROM Win32_CDROMDrive")
		For Each objCD In colCD
			If objCD.MediaLoaded = "True" And objFSO.FileExists(objCD.Drive & "\Sources\boot.wim") Then
				strPEBootDrive = objCD.Drive & "\"
				objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaDrive",strPEBootDrive
				Exit For
			End If
		Next
		
	Else
		strPEBootType = "UNKNOWN"
		objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaType",strPEBootType
		objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,"BootMediaDrive","NONE"
	End If
'==================================================================================================
'WinPE Build Version is automatically calculated by the Last Modified attribute of X:\Windows\System32\Config\System

	Dim f, strYear, strMonth, strDay, strBuildDate, strBuildDateLong
	Set f				= objFSO.GetFile("x:\windows\system32\config\system")
	strYear				= DatePart("yyyy",f.DateLastModified)
	strMonth			= Right(String(2, "0") & Month(f.DateLastModified), 2)
	strDay				= Right(String(2, "0") & Day(f.DateLastModified), 2)
	strBuildDate		= strYear & strMonth & strDay
	strKeyPath			= "SYSTEM\@WinPE"
	strValueName		= "WinPE-BuildDate"
	objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,strValueName,strBuildDate
	strBuildDateLong	= FormatDateTime(f.DateLastModified,vbLongDate)
	strKeyPath			= "SYSTEM\@WinPE"
	strValueName		= "WinPE-BuildDateLong"
	objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,strValueName,strBuildDateLong
'==================================================================================================
'Used for WinPE Expiration.  If booted from PXE the build does not expire
'Requires WinPE Plugin Expire to function
	Dim strBuildExpire, strBuildExpireDays, strBuildExpireCountdown, strBuildExpireDate, strBuildExpireDateLong, strBuildExpireDaysLeft
	If strPEBootType <> "PXE" Then
		strBuildExpire				= "Disabled"	'Set to Enabled for this feature to function, but it is recommended to use the Plugin
		strBuildExpireDays			= 180			'Default Expiration of WinPE.  Doesn't matter unless the strBuildExpire is Enabled
		strBuildExpireCountdown		= 14			'Countdown to give notice to the that WinPE will expire

		'Replace the default values with Plugin Expire
		If RegKeyExists("HKEY_LOCAL_MACHINE\SYSTEM\@WinPE\WinPE-BuildExpire") Then objRegistry.GetStringValue HKEY_LOCAL_MACHINE,"SYSTEM\@WinPE","WinPE-BuildExpire",strBuildExpire
		If RegKeyExists("HKEY_LOCAL_MACHINE\SYSTEM\@WinPE\WinPE-BuildExpireDays") Then objRegistry.GetStringValue HKEY_LOCAL_MACHINE,"SYSTEM\@WinPE","WinPE-BuildExpireDays",strBuildExpireDays
		If RegKeyExists("HKEY_LOCAL_MACHINE\SYSTEM\@WinPE\WinPE-BuildExpireCountdown") Then objRegistry.GetStringValue HKEY_LOCAL_MACHINE,"SYSTEM\@WinPE","WinPE-BuildExpireCountdown",strBuildExpireCountdown
		
		'Convert Strings to Integers
		strBuildExpireDays			= CInt(strBuildExpireDays)
		strBuildExpireCountdown		= CInt(strBuildExpireCountdown)

		'Calculate the proper expiration dates
		strBuildExpireDate			= f.DateLastModified + strBuildExpireDays
			'strValueName			= "WinPE-BuildExpireDate"			'Not working for some reason.  On the ToDo
			'objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,strValueName,strBuildExpireDate
		strBuildExpireDateLong		= FormatDateTime(strBuildExpireDate,vbLongDate)
			strValueName			= "WinPE-BuildExpireDateLong"
			objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,strValueName,strBuildExpireDateLong
		strBuildExpireDaysLeft		= strBuildExpireDate - Date
		strBuildExpireDaysLeft		= Int(strBuildExpireDaysLeft)
			strValueName			= "WinPE-BuildExpireDaysLeft"
			objRegistry.SetStringValue HKEY_LOCAL_MACHINE,strKeyPath,strValueName,strBuildExpireDaysLeft
	End If
'==================================================================================================
	If Lcase(strBuildExpire) = "enabled" Then
	
		If strBuildExpireDaysLeft < strBuildExpireCountdown Then
			MsgBox "This Boot Image was created on " & strBuildDateLong & " and will expire on " & strBuildExpireDateLong & ".  There are " & Round(strBuildExpireDaysLeft,0) & " days remaining.",64,"WinPE Notification"
		ElseIf strBuildExpireDaysLeft < 0 Then
			MsgBox "This Boot Image expired on " & strBuildExpireDateLong & ".  This computer will be shutdown.",64,"WinPE Notification"
			sCmd = "wpeutil.exe shutdown"
			objShell.Run "cmd /c " & sCmd, 1, True
		ElseIf strBuildExpireDaysLeft > strBuildExpireDays Then
			MsgBox "The Date and Time do not appear to be set properly on this computer.  Please update this setting in the BIOS / Firmware.",64,"WinPE Notification"
			sCmd = "wpeutil.exe shutdown"
			objShell.Run "cmd /c " & sCmd, 1, True
		End If
		
	End If
'==================================================================================================
'Launch the WinPE Menu
'Unattend_Generic.xml will only exist if the Unattend.vbs was integrated properly
'This will not run for LiteTouch Boot Images
	If objFSO.FileExists(strMenuDir & "\LaunchMenu.vbs") and objFSO.FileExists(strMenuDir & "\Menu.hta") and objFSO.FileExists(envSystemDrive & "\Unattend_Generic.xml") Then
		sCmd = strMenuDir & "\LaunchMenu.vbs"
		objShell.Run sCmd
		subWinPEMonitor
	Else
		'Check if BGInfo.exe is present and launch if we have the right BGI file
		If objFSO.FileExists(envSystem32 & "\bginfo.exe") Then
			'Priority is given to the BGinfo.bgi file in the Menu directory
			If objFSO.FileExists(strMenuDir & "\BGinfo.bgi") Then
				sCmd = "bginfo " & strMenuDir & "\BGinfo.bgi  /TIMER:0 /NOLICPROMPT"
				objShell.Run sCmd, 2, True
			ElseIf objFSO.FileExists(envSystem32 & "\BGinfo.bgi") Then
				sCmd = "bginfo " & envSystem32 & "\BGinfo.bgi  /TIMER:0 /NOLICPROMPT"
				objShell.Run sCmd, 2, True
			End If
		ElseIf objFSO.FileExists(envSystem32 & "\bginfo64.exe") Then
			If objFSO.FileExists(strMenuDir & "\BGinfo.bgi") Then
				sCmd = "bginfo64 " & strMenuDir & "\BGinfo.bgi  /TIMER:0 /NOLICPROMPT"
				objShell.Run sCmd, 2, True
			ElseIf objFSO.FileExists(envSystem32 & "\BGinfo.bgi") Then
				sCmd = "bginfo64 " & envSystem32 & "\BGinfo.bgi  /TIMER:0 /NOLICPROMPT"
				objShell.Run sCmd, 2, True
			End If
		End If

	'Launch LiteTouch.wsf for Microsoft Deployment
		sCmd = "wscript.exe " & envSystemDrive & "\Deploy\Scripts\LiteTouch.wsf"
		objShell.Run sCmd, 1, True
	End If
'==================================================================================================
'This will keep the WinPE Environment active until a reboot is necessary
Sub subWinPEMonitor
		Dim colProcessList, ProcessCount
		Set colProcessList = objWMIService.ExecQuery ("Select * from Win32_Process Where Name = 'imagex.exe' OR Name = 'dism.exe' OR Name = 'TsManager.exe' OR Name = 'TsmBootstrap.exe' OR Name = 'TSProgressUI.exe'")
		ProcessCount = colProcessList.count

		Do While ProcessCount < 2
			wscript.sleep 10000
			Set colProcessList = objWMIService.ExecQuery ("Select * from Win32_Process Where Name = 'imagex.exe' OR Name = 'dism.exe' OR Name = 'TsManager.exe' OR Name = 'TsmBootstrap.exe' OR Name = 'TSProgressUI.exe'")
			ProcessCount = colProcessList.count
		Loop

		Set colProcessList = objWMIService.ExecQuery ("Select * from Win32_Process Where Name = 'cmd.exe' OR Name = 'imagex.exe' OR Name = 'dism.exe' OR Name = 'TsManager.exe' OR Name = 'TsmBootstrap.exe' OR Name = 'TSProgressUI.exe'")
		ProcessCount = colProcessList.count

		Do While ProcessCount > 0
			wscript.sleep 10000
			Set colProcessList = objWMIService.ExecQuery ("Select * from Win32_Process Where Name = 'cmd.exe' OR Name = 'imagex.exe' OR Name = 'dism.exe' OR Name = 'TsManager.exe' OR Name = 'TsmBootstrap.exe' OR Name = 'TSProgressUI.exe'")
			ProcessCount = colProcessList.count
		Loop

		wscript.sleep 10000
End Sub